from llama_index.readers.database import DatabaseReader
from sqlalchemy import create_engine
import os

# reader = DatabaseReader(
#     scheme="postgresql",
#     host="localhost",
#     port="3306",
#     user=os.getenv("LOCAL_MYSQL_USER"),
#     password=os.getenv("LOCAL_MYSQL_PASSWORD"),
#     dbname="llama-index", # 数据库名称
# )

# 安装数据库依赖  pip install llama_index.readers.database  pip install llama-index sqlalchemy pymysql

db_user = os.getenv("LOCAL_MYSQL_USER")
db_password = os.getenv("LOCAL_MYSQL_PASSWORD")
db_host = "localhost"
db_name = "llama-index"

# 创建SQLAlchemy引擎
engine = create_engine(f"mysql+pymysql://{db_user}:{db_password}@{db_host}/{db_name}")

# 初始化DatabaseReader
reader = DatabaseReader(engine=engine)
print(reader)
print(type(reader.sql_database))


query = "SELECT * FROM eb_user"
documents = reader.load_data(query=query)
print(documents)